import numpy as np
import gzip
import pickle
import time
import os
import torch
from torch.multiprocessing import Pool
from dataclasses import dataclass

from models.gnn_policy import MyGNN
from models.setting import TrainParameters, DEBUG
from branch_and_bound_algorithm.bb_unified import Tree


@dataclass
class MLArgTrain:
    H: np.ndarray = None
    r: np.ndarray = None
    theta: list = None
    file_count: int = None
    prune_policy: str = None
    prune_model: MyGNN = None


class DataCollect(object):
    def __init__(self,
                 train_parameters: TrainParameters = None,
                 train_filepath=None):

        self.train_parameters = train_parameters

        self.filepath = train_filepath

        self.file_count_offset = 0

    def collect_data(self, num_instances):
        # solve the instances of optimization problem, and save the data (features, label) to files.
        N = self.train_parameters.N
        M = self.train_parameters.M
        L = self.train_parameters.L

        H = np.random.randn(num_instances, M, N) + 1j * np.random.randn(num_instances, M, N)
        r = np.random.randn(num_instances, M) + 1j * np.random.randn(num_instances, M)
        theta = np.random.rand(num_instances, N, L) * 2 * np.pi - np.pi # the ranges of the thetas is [-pi,pi]
        theta = np.ndarray.tolist(theta)
        # sort theta[n], n=1,2,...N
        for idx_instance in range(num_instances):
            for row in range(N):
                theta[idx_instance][row].sort()

        arguments_ml = [MLArgTrain(H=H[i], r=r[i], theta=theta[i], file_count=i, prune_policy='default',
                                   prune_model=None) for i in range(len(H))]

        with Pool(min(len(H), 1)) as p:
            out_ml = p.map(self.collect_data_instance, arguments_ml)

        self.file_count_offset += len(H)
        return np.mean([out_ml[i]['non_pruned_node_ratio'] for i in range(len(out_ml))])

    def collect_data_instance(self, arguments: MLArgTrain):

        tree = Tree(epsilon=0.001)
        tree.reset(H=arguments.H,
                   r=arguments.r,
                   theta=arguments.theta,
                   prune_policy=arguments.prune_policy,
                   prune_model=arguments.prune_model)

        t1 = time.time()
        timestep = 0
        time_taken = 0
        num_false_label = 0
        num_label = 0

        while timestep < 1000:
            node_id = tree.select_node()
            if (node_id == None):
                # Find the optimal tree path that includes the optimal solution, and modify the corresponding label
                # to True.
                print('node_id == None')
                print('global_U-global_L/|global_U|:{}'.format(
                    (tree.global_upper_bound - tree.global_lower_bound)
                    / abs(tree.global_upper_bound)))
                node = tree.active_node
                while True:
                    filename = os.path.join(self.filepath, 'sample_{}_{}.pkl'.format(
                        self.file_count_offset + arguments.file_count, node.node_index))

                    with gzip.open(filename, 'rb') as f:
                        (features, _) = pickle.load(f)
                    with gzip.open(filename, 'wb') as f:
                        pickle.dump((features, False), f)
                    num_false_label += 1
                    if node.parent_node:
                        node = node.parent_node
                    else:
                        break
                break
            tree.active_node = tree.nodes[node_id]
            features = tree.get_features()

            # save the features and label(initial value is True) into a file.
            self.save_file((features, True), arguments.file_count, tree.active_node.node_index)
            num_label += 1

            # Add up the calculation time
            time_taken += time.time() - t1
            t1 = time.time()
            branching_var_id = tree.select_branching_variable()
            if (branching_var_id == None):
                # Find the optimal tree path that includes the optimal solution, and modify the corresponding label
                # to True.
                print('branching_var_id == None')
                print('global_U-global_L/|global_U|:{}'.format(
                    (tree.global_upper_bound - tree.global_lower_bound)
                    / abs(tree.global_upper_bound)))
                node = tree.active_node
                while True:
                    filename = os.path.join(self.filepath, 'sample_{}_{}.pkl'.format(
                        self.file_count_offset + arguments.file_count, node.node_index))

                    with gzip.open(filename, 'rb') as f:
                        (features, _) = pickle.load(f)
                    with gzip.open(filename, 'wb') as f:
                        pickle.dump((features, False), f)
                    num_false_label += 1
                    if node.parent_node:
                        node = node.parent_node
                    else:
                        break
                break

            tree.push_children(branching_var_id, node_id)

            # tree.fathom_nodes()
            if (abs(tree.global_upper_bound - tree.global_lower_bound) / abs(tree.global_upper_bound) < tree.epsilon):
                # Find the optimal tree path that includes the optimal solution, and modify the corresponding label
                # to True.
                node = tree.active_node
                while True:
                    filename = os.path.join(self.filepath, 'sample_{}_{}.pkl'.format(
                        self.file_count_offset + arguments.file_count, node.node_index))

                    with gzip.open(filename, 'rb') as f:
                        (features, _) = pickle.load(f)
                    with gzip.open(filename, 'wb') as f:
                        pickle.dump((features, False), f)
                    num_false_label += 1
                    if node.parent_node:
                        node = node.parent_node
                    else:
                        break
                break
            timestep += 1

        return {'time': time_taken, 'non_pruned_node_ratio': num_false_label / num_label}

    def save_file(self, sample, file_count, node_counter):
        if self.filepath is not None:
            filename = os.path.join(self.filepath,
                                    'sample_{}_{}.pkl'.format(self.file_count_offset + file_count, node_counter))
            with gzip.open(filename, 'wb') as f:
                pickle.dump(sample, f)
